"""
用于抽取X-AnyLabeling的图片并进行处理
"""

import os
import shutil

import cv2

# 图片所在的文件夹路径
imagesPath = R"F:\Project\DeepLearning\Game\material\dataset\fourMan\original\material"
# 标注结果的文件夹路径
labeledPath = R"F:\Project\DeepLearning\Game\material\dataset\fourMan\original\labels"
# 绝对路径,需要放入的文件夹地址
purposePath = R"F:\Project\DeepLearning\Game\material\dataset\fourMan"

# 训练集合
trainNames = set()
# 验证集合
validNames = set()
# 测试集合
testNames = set()


#
def readAndDealImages(imagesPath: str):
    """
    获得文件并进行处理
    :param imagesPath: 图片所在的文件夹路径
    """
    allFiles = os.walk(imagesPath)
    # 当前的索引
    count = 0
    # 遍历文件夹内所有的内容
    for root, subFolders, files in allFiles:
        # 获得所有的文件
        for fileName in files:
            # 分离文件名前面的名和后缀
            name, suffix = os.path.splitext(fileName)
            # 绝对路径的文件
            tempAbsFilePath = root + "/" + fileName
            # n抽1,用于作为训练train集
            if count % 10 != 0:
                trainNames.add(name)
                # 复制到指定目录
                shutil.copy(tempAbsFilePath, purposePath + "/train/images/" + fileName)
            # 作为验证valid集合
            else:
                validNames.add(name)
                # 复制到指定目录
                shutil.copy(tempAbsFilePath, purposePath + "/valid/images/" + fileName)
            # elif count % 7 == 0:
            #     validNames.add(name)
            #     # 复制到指定目录
            #     shutil.copy(tempAbsFilePath, purposePath + "/valid/images/" + fileName)
            # # 作为test集合
            # elif count % 17 == 0:
            #     testNames.add(name)
            #     shutil.copy(tempAbsFilePath, purposePath + "/test/images/" + fileName)

            # 自增1
            count += 1


def readAndDealLabels(labeledPath: str):
    """
    获得所有的label并根据图片处理的结果进行转移
    :param labeledPath: label标记好的txt所在的路径
    """
    allFiles = os.walk(labeledPath)
    # 遍历文件夹内所有的文件
    for root, subFolders, files in allFiles:
        # 获得所有的文件
        for fileName in files:
            # 获得后缀
            name, suffix = os.path.splitext(fileName)
            # 只处理txt文件
            if suffix == ".txt":
                # 绝对路径的文件
                tempAbsFilePath = root + "/" + fileName
                if name in trainNames:
                    shutil.copy(tempAbsFilePath, purposePath + "/train/labels/" + fileName)
                elif name in validNames:
                    shutil.copy(tempAbsFilePath, purposePath + "/valid/labels/" + fileName)
                elif name in testNames:
                    shutil.copy(tempAbsFilePath, purposePath + "/test/labels/" + fileName)


if __name__ == '__main__':
    # 清除缓存数据
    trainNames.clear()
    validNames.clear()
    testNames.clear()
    # 查看是否为有效文件夹,内部是否有合理的文件结构
    if not os.path.exists(purposePath):
        exit(-1)
    # 查看是否存在合理的文件结构,没有则建立
    os.makedirs(purposePath + "/train/images", exist_ok=True)
    os.makedirs(purposePath + "/train/labels", exist_ok=True)
    os.makedirs(purposePath + "/valid/images", exist_ok=True)
    os.makedirs(purposePath + "/valid/labels", exist_ok=True)
    os.makedirs(purposePath + "/test/images", exist_ok=True)
    os.makedirs(purposePath + "/test/labels", exist_ok=True)

    # 读取图片并处理
    readAndDealImages(imagesPath)
    # 根据处理的图片,处理label
    readAndDealLabels(labeledPath)
